import os
import random
import torch
import argparse
import warnings

import torch.backends.cudnn as cudnn

parser = argparse.ArgumentParser(description='SimKGC arguments')
#--------end------------SimKGC------------的参数
#uncased不能区分大小写
parser.add_argument('--pretrained-model', default='hfl/chinese-bert-wwm', type=str, metavar='N',
                    help='path to pretrained;model:hfl/chinese-bert-wwm,bert-base-chinese,bert-base-uncased,')
parser.add_argument('--task', default='MEDsmall', type=str, metavar='N',
                    help='dataset name:wn18rr,')
parser.add_argument('--train-path', default='', type=str, metavar='N',
                    help='path to training data')
parser.add_argument('--valid_path', default='', type=str, metavar='N',
                    help='path to valid data')
parser.add_argument('--model-dir', default='', type=str, metavar='N',
                    help='path to model dir')
#学习率预热的方法，它在训练开始的时候先选择使用一个较小的学习率,后恢复正常
parser.add_argument('--warmup', default=400, type=int, metavar='N',
                    help='warmup steps')
parser.add_argument('--max-to-keep', default=5, type=int, metavar='N',
                    help='max number of checkpoints to keep')
parser.add_argument('--grad-clip', default=10.0, type=float, metavar='N',
                    help='gradient clipping')
parser.add_argument('--pooling', default='cls', type=str, metavar='N',
                    help='bert pooling')
parser.add_argument('--dropout', default=0.1, type=float, metavar='N',
                    help='dropout on final linear layer')
##梯度缩放，防止参数下溢
parser.add_argument('--use-amp', action='store_true',
                    help='Use amp if available')
parser.add_argument('--t', default=0.05, type=float,
                    help='temperature parameter')
parser.add_argument('--use-link-graph', action='store_true',
                    help='use neighbors from link graph as context')
parser.add_argument('--eval-every-n-step', default=10000, type=int,
                    help='evaluate every n steps')
parser.add_argument('--pre-batch', default=0, type=int,
                    help='number of pre-batch used for negatives')
parser.add_argument('--pre-batch-weight', default=0.5, type=float,
                    help='the weight for logits from pre-batch negatives')
parser.add_argument('--additive-margin', default=0.0, type=float, metavar='N',
                    help='additive margin for InfoNCE loss function')#γ
parser.add_argument('--finetune-t', action='store_true',
                    help='make temperature as a trainable parameter or not')
parser.add_argument('--max-num-tokens', default=50, type=int,
                    help='maximum number of tokens')
#action='store_true' 是指带触发action时为真(运行时带此参数则为真)，不触发则为假。
parser.add_argument('--use-self-negative', action='store_true',
                    help='use head entity as negative')

parser.add_argument('-j', '--workers', default=1, type=int, metavar='N',
                    help='number of data loading workers')
parser.add_argument('--epochs',default=10,dest='max_epochs',  type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('-b', '--batch_size', default=128, type=int,
                    metavar='N',
                    help='mini-batch size (default: 256), this is the total '
                         'batch size of all GPUs on the current node when '
                         'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr', '--learning-rate', default=2e-5, type=float,
                    metavar='LR', help='initial learning rate', dest='lr')
parser.add_argument('--lr-scheduler', default='linear', type=str,
                    help='Lr scheduler to use')
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)',
                    dest='weight_decay')
parser.add_argument('-p', '--print-freq', default=50, type=int,
                    metavar='N', help='print frequency (default: 10)')
# parser.add_argument('--seed', default=None, type=int,help='seed for initializing training. ')

# only used for evaluation
parser.add_argument('--is-test', action='store_true',
                    help='is in test mode or not')
parser.add_argument('--rerank-n-hop', default=2, type=int,
                    help='use n-hops node for re-ranking entities, only used during evaluation')
parser.add_argument('--neighbor-weight', default=0.0, type=float,
                    help='weight for re-ranking entities')
parser.add_argument('--eval-model-path', default='', type=str, metavar='N',
                    help='path to model, only used for evaluation')
#--------end------------SimKGC-----------的参数-

#以下是来自compgcn的参数
parser.add_argument('-name',		default='testrun',					help='Set run name for saving/restoring models')
parser.add_argument('-data',		dest='dataset',         default='',            help='FB15k-237/WN18RR')
parser.add_argument('-model',		dest='model',		default='compgcn',		help='or ')
parser.add_argument('-score_func',	dest='score_func',	default='conve',		help='Score Function for Link prediction')
parser.add_argument('-opn',             dest='opn',             default='corr',                 help='Composition Operation to be used in CompGCN')

# parser.add_argument('-batch',           dest='batch_size',      default=128,    type=int,       help='Batch size')
parser.add_argument('-gamma',		type=float,             default=40.0,			help='Margin')
parser.add_argument('-gpu',		type=str,               default='0',			help='Set GPU Ids : Eg: For CPU = -1, For Single GPU = 0')
# parser.add_argument('-epoch',		dest='max_epochs', 	type=int,       default=500,  	help='Number of epochs')
parser.add_argument('-l2',		type=float,             default=0.0,			help='L2 Regularization for Optimizer')
# parser.add_argument('-lr',		type=float,             default=0.001,			help='Starting Learning Rate')
parser.add_argument('-lbl_smooth',      dest='lbl_smooth',	type=float,     default=0.1,	help='Label Smoothing')
# parser.add_argument('-num_workers',	type=int,               default=10,                     help='Number of processes to construct batches')
parser.add_argument('-seed',            dest='seed',            default=41504,  type=int,     	help='Seed for randomization')

parser.add_argument('-restore',         dest='restore',         action='store_true',            help='Restore from the previously saved model')
parser.add_argument('-bias',            dest='bias',            action='store_true',            help='Whether to use bias in the model')

parser.add_argument('-num_bases',	dest='num_bases', 	default=-1,   	type=int, 	help='Number of basis relation vectors to use')
parser.add_argument('-init_dim',	dest='init_dim',	default=100,	type=int,	help='Initial dimension size for entities and relations')
parser.add_argument('-gcn_dim',	  	dest='gcn_dim', 	default=200,   	type=int, 	help='Number of hidden units in GCN')
parser.add_argument('-embed_dim',	dest='embed_dim', 	default=None,   type=int, 	help='Embedding dimension to give as input to score function')
parser.add_argument('-gcn_layer',	dest='gcn_layer', 	default=1,   	type=int, 	help='Number of GCN Layers to use')
parser.add_argument('-gcn_drop',	dest='dropout', 	default=0.1,  	type=float,	help='Dropout to use in GCN Layer')
parser.add_argument('-hid_drop',  	dest='hid_drop', 	default=0.3,  	type=float,	help='Dropout after GCN')

# ConvE specific hyperparameters
parser.add_argument('-hid_drop2',  	dest='hid_drop2', 	default=0.3,  	type=float,	help='ConvE: Hidden dropout')
parser.add_argument('-feat_drop', 	dest='feat_drop', 	default=0.3,  	type=float,	help='ConvE: Feature Dropout')
parser.add_argument('-k_w',	  	dest='k_w', 		default=10,   	type=int, 	help='ConvE: k_w')
parser.add_argument('-k_h',	  	dest='k_h', 		default=20,   	type=int, 	help='ConvE: k_h')
parser.add_argument('-num_filt',  	dest='num_filt', 	default=200,   	type=int, 	help='ConvE: Number of filters in convolution')
parser.add_argument('-ker_sz',    	dest='ker_sz', 		default=7,   	type=int, 	help='ConvE: Kernel size to use')

parser.add_argument('-logdir',          dest='log_dir',         default='./log/',               help='Log directory')
parser.add_argument('-config',          dest='config_dir',      default='./config/',            help='Config directory')
#以上是来自compgcn的参数

#-----------------start-----GCN_rel.gcn_t.py的参数
parser.add_argument('--device',type=str,default='cuda:0',help='')
# parser.add_argument('--data',type=str,default='data/2022waiting',help='data path') #去富川！#28个之前是Bucheon文件夹
parser.add_argument('--adjdata',type=str,default='gcn_model/data/sensor_graph/adj.pkl',help='adj data path') #实际邻接.
parser.add_argument('--adjtype',type=str,default='doubletransition',help='adj type')
parser.add_argument('--gcn_bool',action='store_true',help='whether to add graph convolution layer')
parser.add_argument('--aptonly',action='store_true',help='whether only adaptive adj')
parser.add_argument('--addaptadj',action='store_true',help='whether add adaptive adj')
parser.add_argument('--randomadj',action='store_true',help='whether random initialize adaptive adj')#随机初始化A
parser.add_argument('--seq_length',type=int,default=12,help='')
parser.add_argument('--nhid',type=int,default=32,help='')
parser.add_argument('--in_dim',type=int,default=2,help='inputs dimension') #2
parser.add_argument('--num_nodes',type=int,default=572,help='number of nodes') #214
# parser.add_argument('--batch_size',type=int,default=64,help='batch size')
parser.add_argument('--learning_rate',type=float,default=0.001,help='learning rate')
parser.add_argument('--gcn_dropout',type=float,default=0.3,help='dropout rate')
parser.add_argument('--gcn_layer_num',type=int,default=2,help='order')
# parser.add_argument('--weight_decay',type=float,default=0.0001,help='weight decay rate')
# parser.add_argument('--epochs',type=int,default=50,help='') #
parser.add_argument('--print_every',type=int,default=50,help='')
#parser.add_argument('--seed',type=int,default=99,help='random seed')
parser.add_argument('--save',type=str,default='./train_result_path/waiting',help='save path')
parser.add_argument('--expid',type=int,default=1,help='experiment id')
parser.add_argument('--c_in',type=int,default=1,help='gcn_model in_dim')
parser.add_argument('--c_out',type=int,default=1,help='gcn_model out_dim')
#-----------------end-----GCN_rel.gcn_t.py的参数

args = parser.parse_args()

# assert not args.train_path or os.path.exists(args.train_path)
# assert args.pooling in ['cls', 'mean', 'max']
# # assert args.task.lower() in ['wn18rr', 'fb15k237', 'wiki5m_ind', 'wiki5m_trans']
# assert args.task.lower() in ['wn18rr', 'fb15k237', 'wiki5m_ind', 'wiki5m_trans','own_disease','med_body_clin','med']
# assert args.lr_scheduler in ['linear', 'cosine']

if args.model_dir:
    os.makedirs(args.model_dir, exist_ok=True)
else:
    # assert os.path.exists(args.eval_model_path), 'One of args.model_dir and args.eval_model_path should be valid path'
    args.model_dir = os.path.dirname(args.eval_model_path)

if args.seed is not None:
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    cudnn.deterministic = True

try:
    if args.use_amp:
        import torch.cuda.amp
except Exception:
    args.use_amp = False
    warnings.warn('AMP training is not available, set use_amp=False')

if not torch.cuda.is_available():
    args.use_amp = False
    args.print_freq = 1
    warnings.warn('GPU is not available, set use_amp=False and print_freq=1')

